{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recommender Systems with GNNs\n",
    "\n",
    "The academia and industries have witnessed various attempts and applications of applying GNNs to business applications including recommender systems, fraud detection, etc.\n",
    "\n",
    "In this hands-on tutorial, we will see how GNNs can be used for a recommender system.\n",
    "\n",
    "## Recommender System: Overview\n",
    "\n",
    "Recommender systems are typically used on E-commerce or similar platforms.  It is responsible for predicting the items that are most likely to be interacted by a user, given his/her previous interaction history.\n",
    "\n",
    "Based on whether the interactions have a *rating* that signifies how a user likes an item, we could categorize the interaction data as having either *explicit feedback* if the data has ratings, or *implicit feedback* otherwise.\n",
    "\n",
    "There are a lot of methodologies for recommendation.  Here we will focus on *matrix-factorization-based* methods that expresses users and items as *embeddings* (i.e. vectors of same dimensionality), and the preference of a user to an item as the *dot product* between user and item embeddings.\n",
    "\n",
    "More specifically, for each user we can assign a learnable vector $\\boldsymbol{u}_i$, and for each item we can assign another learnable vector $\\boldsymbol{v}_j$.  Additionally, we also learn a bias for each user $\\beta_i$ and each item $\\gamma_j$, representing the baseline rating of each user and item.  Preference prediction is expressed as $\\hat{r}_{i,j}=\\boldsymbol{u}_i^\\top \\boldsymbol{v}_j + \\beta_i + \\gamma_j$.\n",
    "\n",
    "For explicit feedback datasets, we try to minimize the difference between prediction and ground truth $r_{i,j}$ for all user-item interactions:\n",
    "\n",
    "$$\n",
    "\\min_{\\boldsymbol{u},\\boldsymbol{v}} \\sum_{(i,j)\\in \\mathcal{D}} \\left(\\hat{r}_{i,j} - r_{i,j}\\right)^2\n",
    "$$\n",
    "\n",
    "For implicit feedback datasets, the function to minimize is more complicated, so in this tutorial we will focus on explicit feedback datasets.\n",
    "\n",
    "Other recommender system families include neighborhood-based models, factorization machines, etc.  But we won't go through them here.\n",
    "\n",
    "## What Does a GNN Do in Recommender Systems?\n",
    "\n",
    "GNNs most commonly appear as a another method of computing the user and item embeddings: instead of directly assigning a learnable vector as user and item embedding, we compute the embeddings with a GNN.  The graph for the GNN would be the interaction data itself, with users and items as two types of nodes and the interactions as edges connecting the nodes.  \n",
    "\n",
    "Comparing against other approaches, the biggest advantage of GNNs in recommender systems is that it can potentially combine the interactions data with other kind of relational data.  If we have additional relational data such as social networks or knowledge graphs, we can combine those data with the interaction data to form a bigger graph, and still run the same GNN as a recommender system.  Other advantages include ability of introducing the information of neighbors to embeddings, and the ability to combine both user/item features and the neighborhood structure of users/items, etc.\n",
    "\n",
    "In our hands-on, we will use the MovieLens-100K dataset.\n",
    "\n",
    "Note that to simplify stuff in our tutorial, we would not be considering new users and items."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "train_data = pd.read_csv('ua.base', sep='\\t', header=None, names=['user_id', 'item_id', 'rating', 'timestamp'])\n",
    "test_data = pd.read_csv('ua.test', sep='\\t', header=None, names=['user_id', 'item_id', 'rating', 'timestamp'])\n",
    "user_data = pd.read_csv('u.user', sep='|', header=None, encoding='latin1')\n",
    "item_data = pd.read_csv('u.item', sep='|', header=None, encoding='latin1')\n",
    "\n",
    "# Remove the entries with users and items not appearing in the training dataset\n",
    "test_data = test_data[test_data['user_id'].isin(train_data['user_id']) &\n",
    "                      test_data['item_id'].isin(train_data['item_id'])]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is a preview of the tables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_data.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "item_data.head(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Heterogeneous Graphs in DGL\n",
    "\n",
    "If we treat our interaction data as a graph, we will have two types of nodes: users and items.  We call a graph that has multiple node types and edge types a *heterogeneous* graph.\n",
    "\n",
    "With MovieLens-100K loaded, we will see how to build a heterogeneous graph from it.  First, we need to figure out what are the nodes and edges.\n",
    "\n",
    "Here, since we have a user-item interaction dataset, we simply designate users and items as nodes, and interactions as edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dgl\n",
    "import torch\n",
    "\n",
    "# Relabel the user IDs and item IDs to consecutive integers\n",
    "train_data = train_data.astype({'user_id': 'category', 'item_id': 'category'})\n",
    "test_data = test_data.astype({'user_id': 'category', 'item_id': 'category'})\n",
    "# We need to keep the relabeling between training and test data consistent\n",
    "test_data['user_id'].cat.set_categories(train_data['user_id'].cat.categories, inplace=True)\n",
    "test_data['item_id'].cat.set_categories(train_data['item_id'].cat.categories, inplace=True)\n",
    "\n",
    "train_user_ids = torch.LongTensor(train_data['user_id'].cat.codes.values)\n",
    "train_item_ids = torch.LongTensor(train_data['item_id'].cat.codes.values)\n",
    "train_ratings = torch.LongTensor(train_data['rating'].values)\n",
    "test_user_ids = torch.LongTensor(test_data['user_id'].cat.codes.values)\n",
    "test_item_ids = torch.LongTensor(test_data['item_id'].cat.codes.values)\n",
    "test_ratings = torch.LongTensor(test_data['rating'].values)\n",
    "\n",
    "# Build graph\n",
    "graph = dgl.heterograph({\n",
    "    # Heterogeneous graphs are organized as a dictionary of edges connecting two types of nodes.\n",
    "    # We specify the edges of a type simply with a pair of user ID array and item ID array.\n",
    "    ('user', 'watched', 'item'): (train_user_ids, train_item_ids),\n",
    "    # Since DGL graphs are directional, we need an inverse relation from items to users as well.\n",
    "    ('item', 'watched-by', 'user'): (train_item_ids, train_user_ids)\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can visualize the \"metagraph\" of the heterogeneous graph.\n",
    "\n",
    "A *metagraph* in DGL describes how edge types connect with different node types.  It is equivalent to the schema of a heterogeneous information network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import networkx as nx\n",
    "\n",
    "# The metagraph is shown in the figure below.  Note the bidirectional arrows.\n",
    "nx.draw_networkx(graph.metagraph, node_color='white')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we find which users and items appeared in training and test set, we now need to assign corresponding node features to the user and item nodes.  For simplicity, we would only consider the following features:\n",
    "\n",
    "* Movie genres: the movie genres are expressed as 0-1 masks in the MovieLens-100K dataset.  We will take the mask as feature vector of the movie.\n",
    "* User age: we will bin the user age by decades and treat the bin as a categorical variable.\n",
    "* User gender: a categorical variable.\n",
    "* User occupation: a categorical variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_data[0] = user_data[0].astype('category')\n",
    "user_data[0] = user_data[0].cat.set_categories(train_data['user_id'].cat.categories)\n",
    "user_data = user_data.dropna(subset=[0])\n",
    "user_data[0] = user_data[0].cat.codes\n",
    "user_data = user_data.sort_values(0)\n",
    "item_data[0] = item_data[0].astype('category')\n",
    "item_data[0] = item_data[0].cat.set_categories(train_data['item_id'].cat.categories)\n",
    "item_data = item_data.dropna(subset=[0])\n",
    "item_data[0] = item_data[0].cat.codes\n",
    "item_data = item_data.sort_values(0)\n",
    "\n",
    "# Convert the age, gender, and occupation column to categorical\n",
    "user_data[2] = user_data[2].astype('category')\n",
    "user_data[3] = user_data[3].astype('category')\n",
    "user_data[4] = user_data[4].astype('category')\n",
    "\n",
    "user_age = user_data[1].values // 10\n",
    "num_user_age_bins = user_age.max() + 1     # count the number of user age bins\n",
    "user_gender = user_data[2].cat.codes.values\n",
    "num_user_genders = len(user_data[2].cat.categories)\n",
    "user_occupation = user_data[3].cat.codes.values\n",
    "num_user_occupations = len(user_data[3].cat.categories)\n",
    "\n",
    "item_genres = item_data[range(5, 24)].values\n",
    "num_item_genres = item_genres.shape[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Manipulating node and edge features can be accomplished with the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assign user features\n",
    "graph.nodes['user'].data['age'] = torch.LongTensor(user_age)\n",
    "graph.nodes['user'].data['gender'] = torch.LongTensor(user_gender)\n",
    "graph.nodes['user'].data['occupation'] = torch.LongTensor(user_occupation)\n",
    "# Assign item features\n",
    "graph.nodes['item'].data['genres'] = torch.FloatTensor(item_genres)\n",
    "# Assign ratings\n",
    "graph.edges['watched'].data['rating'] = torch.LongTensor(train_ratings)\n",
    "graph.edges['watched-by'].data['rating'] = torch.LongTensor(train_ratings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model on Heterogeneous Graphs\n",
    "\n",
    "### Define Dataset to Sample Minibatches from\n",
    "\n",
    "The first thing we need to do for GNN training is to define how a minibatch of examples look like.  For rating prediction tasks, a minibatch consists of pairs of users and items.  So we create a torch `Dataset` object that contains all training pairs of users and items as well as their ratings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "\n",
    "train_dataset = TensorDataset(train_user_ids, train_item_ids, train_ratings)\n",
    "test_dataset = TensorDataset(test_user_ids, test_item_ids, test_ratings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Minibatch & Neighbor Sampler\n",
    "\n",
    "For an interaction we wish to compute the representation of its user and its item with a multi-layer GNN.  Therefore, as in the previous tutorial, we need to define a minibatch sampler that gets us the computation dependency of the user and item nodes for the multi-layer GNN.  At the same time, we also want to keep track of the training interactions themselves, i.e. which user is interacting which item.\n",
    "\n",
    "In DGL, rating prediction could be achieved with the following steps:\n",
    "\n",
    "1. Create a *pair graph* that consists of all the users and items in the dataset, but only with the interactions in the training minibatch.\n",
    "2. *Compact* the graph to remove all unnecessary nodes, keeping only the nodes with some edges going in or out.  We provide a function `dgl.compact_graphs` for this.\n",
    "3. Construct blocks bottom-up for computing the multi-layer output of the nodes, as in the previous tutorial.\n",
    "4. Propagate messages top-down to obtain the output from the GNN.\n",
    "5. Copy the output to the pair graph and compute the preference with `apply_edges()` method.\n",
    "\n",
    "The following cell has the code of completing step 1 to step 3, which constitutes the steps of obtaining computation dependency for the interactions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MinibatchSampler(object):\n",
    "    def __init__(self, graph, num_layers):\n",
    "        self.graph = graph\n",
    "        self.num_layers = num_layers\n",
    "        \n",
    "    def sample(self, batch):\n",
    "        # Convert the list of user-item-rating triplets into a triplet of users, items, and ratings\n",
    "        users, items, ratings = zip(*batch)\n",
    "        users = torch.stack(users)\n",
    "        items = torch.stack(items)\n",
    "        ratings = torch.stack(ratings)\n",
    "        \n",
    "        # Create a pair graph (Step 1)\n",
    "        pair_graph = dgl.heterograph(\n",
    "            {('user', 'watched', 'item'): (users, items)},\n",
    "            num_nodes_dict={'user': self.graph.number_of_nodes('user'), 'item': self.graph.number_of_nodes('item')})\n",
    "        \n",
    "        # Compact the graph (Step 2)\n",
    "        pair_graph = dgl.compact_graphs(pair_graph)\n",
    "        # Assign ratings to the graph\n",
    "        pair_graph.edata['rating'] = ratings\n",
    "        \n",
    "        # Construct blocks (Step 3)\n",
    "        seeds = {'user': pair_graph.nodes['user'].data[dgl.NID],\n",
    "                 'item': pair_graph.nodes['item'].data[dgl.NID]}\n",
    "        blocks = self.construct_blocks(seeds, (users, items))\n",
    "        \n",
    "        # Copy node features from original graph to the sampled block.\n",
    "        # Note that for our model we only need to copy the features to the source side of the first block.\n",
    "        # The node features of other blocks would be computed by our model.\n",
    "        for feature_name in self.graph.nodes['user'].data.keys():\n",
    "            blocks[0].srcnodes['user'].data[feature_name] = \\\n",
    "                self.graph.nodes['user'].data[feature_name][blocks[0].srcnodes['user'].data[dgl.NID]]\n",
    "        for feature_name in self.graph.nodes['item'].data.keys():\n",
    "            blocks[0].srcnodes['item'].data[feature_name] = \\\n",
    "                self.graph.nodes['item'].data[feature_name][blocks[0].srcnodes['item'].data[dgl.NID]]\n",
    "\n",
    "        return pair_graph, blocks\n",
    "        \n",
    "    def construct_blocks(self, seeds, user_item_pairs_to_remove):\n",
    "        blocks = []\n",
    "        users, items = user_item_pairs_to_remove\n",
    "        for i in range(self.num_layers):\n",
    "            # We take all neighbors to form the sampled graph for computing the node representations on the\n",
    "            # current layer.\n",
    "            sampled_graph = dgl.in_subgraph(self.graph, seeds)\n",
    "            # Find the sampled edge IDs for both directions\n",
    "            sampled_eids = sampled_graph.edges['watched'].data[dgl.EID]\n",
    "            sampled_eids_rev = sampled_graph.edges['watched-by'].data[dgl.EID]\n",
    "            \n",
    "            # A subtlety of rating prediction and link prediction is that, when we train on the pair of user A\n",
    "            # and item 1, we don't want to actually tell the GNN that \"user A has a connection to item 1\".  So\n",
    "            # we should remove all edges connecting the training pairs from the sampled graph.\n",
    "            _, _, edges_to_remove = sampled_graph.edge_ids(users, items, etype='watched', return_uv=True)\n",
    "            _, _, edges_to_remove_rev = sampled_graph.edge_ids(items, users, etype='watched-by', return_uv=True)\n",
    "            sampled_with_edges_removed = dgl.remove_edges(\n",
    "                sampled_graph, {'watched': edges_to_remove, 'watched-by': edges_to_remove_rev})\n",
    "            sampled_eids = sampled_eids[sampled_with_edges_removed.edges['watched'].data[dgl.EID]]\n",
    "            sampled_eids_rev = sampled_eids_rev[sampled_with_edges_removed.edges['watched-by'].data[dgl.EID]]\n",
    "            \n",
    "            # Create a block from the sampled graph.\n",
    "            block = dgl.to_block(sampled_with_edges_removed, seeds)\n",
    "            blocks.insert(0, block)\n",
    "            seeds = {'user': block.srcnodes['user'].data[dgl.NID],\n",
    "                     'item': block.srcnodes['item'].data[dgl.NID]}\n",
    "            \n",
    "            # Copy the ratings to the edges of the sampled block\n",
    "            block.edges['watched'].data['rating'] = \\\n",
    "                self.graph.edges['watched'].data['rating'][sampled_eids]\n",
    "            block.edges['watched-by'].data['rating'] = \\\n",
    "                self.graph.edges['watched-by'].data['rating'][sampled_eids_rev]\n",
    "            \n",
    "        return blocks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Model\n",
    "\n",
    "We choose a simplified version of Graph Convolutional Matrix Completion (GCMC) as the model.  Each node sends a message to its neighbors by computing a rating-specific projection of the node's representation.  Then each node gathers the messages and takes an average, projecting it afterwards.\n",
    "\n",
    "More specifically, given a user node $i$ with its input representations $u_i^{l-1}$ on layer $l$, we\n",
    "\n",
    "1. Find all neighbors, i.e. interacted items, of $i$.\n",
    "2. For each item $j$, check the rating $r_{ij}$ between user $i$ and item $j$.  The message sent from item $j$ to user $i$ would be a rating-specific linear projection $m_{j\\to i}^l \\gets W_{r_{ij}}^l v_j^{l-1}$.  $W_{r_{ij}}^l$ is a learnable matrix.\n",
    "3. The user $i$ aggregates all incoming messages by summing and updates itself: $u_i^l \\gets \\mathrm{ReLU}(W^l[\\sum(m_{j\\to i}^l); u_i^{l-1}])$.  $W^l$ is again a learnable matrix.\n",
    "\n",
    "We can similarly compute $v_j$, the representation of an item node $j$.\n",
    "\n",
    "For an implementation faithful to the paper, please see our model examples.\n",
    "\n",
    "#### Implementation in DGL\n",
    "\n",
    "Here we are implementing a *heterogeneous graph convolution*.  It can be decomposed into the following steps:\n",
    "\n",
    "1. Implement or specify the convolution along each edge type individually.\n",
    "2. Use `dgl.nn.HeteroGraphConv` to perform convolution on each edge type individually and aggregate the result of each node type along its incoming edge types.\n",
    "\n",
    "The following is an implementation of GCMC convolution on a single edge type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "import dgl.function as fn\n",
    "import dgl.nn as dglnn\n",
    "\n",
    "class GCMCConv(nn.Module):\n",
    "    def __init__(self, hidden_dims, num_ratings):\n",
    "        super().__init__()\n",
    "        \n",
    "        # The ratings are ranged from 1 to num_ratings, so I add 1 to the number of parameters.\n",
    "        self.W_r = nn.Parameter(torch.randn(num_ratings + 1, hidden_dims, hidden_dims))\n",
    "        self.W = nn.Linear(hidden_dims * 2, hidden_dims)\n",
    "        \n",
    "    def compute_message(self, W, edges):\n",
    "        W_r = W[edges.data['rating']]\n",
    "        h = edges.src['h']\n",
    "        m = (W_r @ h.unsqueeze(-1)).squeeze(2)\n",
    "        return m\n",
    "    \n",
    "    def forward(self, graph, node_features):\n",
    "        with graph.local_scope():\n",
    "            src_features, dst_features = node_features\n",
    "            graph.srcdata['h'] = src_features\n",
    "            graph.dstdata['h'] = dst_features\n",
    "            # Compute messages\n",
    "            graph.apply_edges(lambda edges: {'m': self.compute_message(self.W_r, edges)})\n",
    "            # Aggregate messages\n",
    "            graph.update_all(fn.copy_e('m', 'm'), fn.mean('m', 'h_neigh'))\n",
    "            # Updates the representations of output users and items\n",
    "            result = F.relu(self.W(torch.cat([graph.dstdata['h'], graph.dstdata['h_neigh']], 1)))\n",
    "            return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This object performs a heterogeneous graph convolution on both the user-watched-item relation and item-watched-by-user relation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GCMCLayer(nn.Module):\n",
    "    def __init__(self, hidden_dims, num_ratings):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.heteroconv = dglnn.HeteroGraphConv(\n",
    "            # NN Module for message passing on each individual edge type\n",
    "            {'watched': GCMCConv(hidden_dims, num_ratings), 'watched-by': GCMCConv(hidden_dims, num_ratings)},\n",
    "            # Aggregation strategy on each node type along incoming edge types.\n",
    "            aggregate='sum')\n",
    "        \n",
    "    def forward(self, block, input_user_features, input_item_features):\n",
    "        with block.local_scope():\n",
    "            # Start from the embedding of users and items...\n",
    "            h_user = input_user_features\n",
    "            h_item = input_item_features\n",
    "            \n",
    "            src_features = {'user': h_user, 'item': h_item}\n",
    "            # First copy the features from the source side to the destination side.\n",
    "            # The first few nodes of the source node side would be identical to that of the destination\n",
    "            # side, so we can just do the following:\n",
    "            dst_features = {'user': h_user[:block.number_of_dst_nodes('user')], 'item': h_item[:block.number_of_dst_nodes('item')]}\n",
    "            \n",
    "            # HeteroGraphConv essentially performs a heterogeneous graph convolution by:\n",
    "            # (1) Performing message passing on each edge type individually.  In this case,\n",
    "            #     * Items receives messages from users along \"watched\" relation.\n",
    "            #     * Users receives messages from items along \"watched-by\" relation.\n",
    "            # (2) Aggregating the result on each node type along incoming edge types.  In this case,\n",
    "            #     * Nothing happens since both users and items only have one incoming edge type.\n",
    "            result = self.heteroconv(block, (src_features, dst_features))\n",
    "            return result['user'], result['item']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally this object handles both the message passing and score computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GCMCRating(nn.Module):\n",
    "    def __init__(self, num_users, num_items, hidden_dims, num_ratings, num_layers):\n",
    "        super().__init__()\n",
    "        \n",
    "        # Node-specific learnable embeddings\n",
    "        self.user_embeddings = nn.Embedding(num_users, hidden_dims)\n",
    "        self.item_embeddings = nn.Embedding(num_items, hidden_dims)\n",
    "        \n",
    "        # Transformation modules for input features of users and items\n",
    "        self.U_age = nn.Embedding(num_user_age_bins, hidden_dims)\n",
    "        self.U_gender = nn.Embedding(num_user_genders, hidden_dims)\n",
    "        self.U_occupation = nn.Embedding(num_user_occupations, hidden_dims)\n",
    "        self.U_genres = nn.Linear(num_item_genres, hidden_dims)\n",
    "        \n",
    "        self.layers = nn.ModuleList([\n",
    "            GCMCLayer(hidden_dims, num_ratings) for _ in range(num_layers)])\n",
    "        \n",
    "        self.W = nn.Linear(hidden_dims, hidden_dims)\n",
    "        self.V = nn.Linear(hidden_dims, hidden_dims)\n",
    "        \n",
    "    def forward(self, blocks):\n",
    "        # Propagate messages top-down (Step 4)\n",
    "        # We start with a learnable embedding for each user and item...\n",
    "        user_embeddings = self.user_embeddings(blocks[0].srcnodes['user'].data[dgl.NID])\n",
    "        item_embeddings = self.item_embeddings(blocks[0].srcnodes['item'].data[dgl.NID])\n",
    "        \n",
    "        # And add the transformation from corresponding user and item features..\n",
    "        user_embeddings = user_embeddings + self.U_age(blocks[0].srcnodes['user'].data['age'])\n",
    "        user_embeddings = user_embeddings + self.U_gender(blocks[0].srcnodes['user'].data['gender'])\n",
    "        user_embeddings = user_embeddings + self.U_occupation(blocks[0].srcnodes['user'].data['occupation'])\n",
    "        item_embeddings = item_embeddings + self.U_genres(blocks[0].srcnodes['item'].data['genres'])\n",
    "        \n",
    "        # Then perform a heterogeneous GCMC convolution\n",
    "        for block, layer in zip(blocks, self.layers):\n",
    "            user_embeddings, item_embeddings = layer(block, user_embeddings, item_embeddings)\n",
    "        \n",
    "        # Compute predicted preference (Step 5)\n",
    "        user_embeddings = self.W(user_embeddings)\n",
    "        item_embeddings = self.V(item_embeddings)\n",
    "        \n",
    "        return user_embeddings, item_embeddings\n",
    "        \n",
    "    def compute_score(self, pair_graph, user_embeddings, item_embeddings):\n",
    "        with pair_graph.local_scope():\n",
    "            pair_graph.nodes['user'].data['h'] = user_embeddings\n",
    "            pair_graph.nodes['item'].data['h'] = item_embeddings\n",
    "            pair_graph.apply_edges(fn.u_dot_v('h', 'h', 'r'))\n",
    "            \n",
    "            return pair_graph.edata['r']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Evaluation Metric\n",
    "\n",
    "In this tutorial we use root of mean squared error (RMSE) as the metric."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rmse(pred, label):\n",
    "    return ((pred - label) ** 2).mean().sqrt()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Loop\n",
    "\n",
    "Note: in this tutorial we are directly validating on the test set for simplicity.  In practice we should split out some data points in the training set to form a validation set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm\n",
    "\n",
    "# In this tutorial we consider 1-layer GNNs.\n",
    "NUM_LAYERS = 1\n",
    "BATCH_SIZE = 500\n",
    "NUM_EPOCHS = 50\n",
    "HIDDEN_DIMS = 8\n",
    "\n",
    "sampler = MinibatchSampler(graph, NUM_LAYERS)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=BATCH_SIZE, collate_fn=sampler.sample, shuffle=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=BATCH_SIZE, collate_fn=sampler.sample, shuffle=False)\n",
    "\n",
    "model = GCMCRating(graph.number_of_nodes('user'), graph.number_of_nodes('item'), HIDDEN_DIMS, 5, NUM_LAYERS)\n",
    "opt = torch.optim.Adam(model.parameters())\n",
    "\n",
    "for _ in range(NUM_EPOCHS):\n",
    "    model.train()\n",
    "    with tqdm.tqdm(train_dataloader) as t:\n",
    "        for pair_graph, blocks in t:\n",
    "            user_emb, item_emb = model(blocks)\n",
    "            prediction = model.compute_score(pair_graph, user_emb, item_emb)\n",
    "            loss = ((prediction - pair_graph.edata['rating']) ** 2).mean()\n",
    "            opt.zero_grad()\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "            t.set_postfix({'loss': '%.4f' % loss.item()}, refresh=False)\n",
    "    model.eval()\n",
    "    with tqdm.tqdm(test_dataloader) as t:\n",
    "        with torch.no_grad():\n",
    "            predictions = []\n",
    "            ratings = []\n",
    "            for pair_graph, blocks in t:\n",
    "                user_emb, item_emb = model(blocks)\n",
    "                prediction = model.compute_score(pair_graph, user_emb, item_emb)\n",
    "                predictions.append(prediction)\n",
    "                ratings.append(pair_graph.edata['rating'])\n",
    "\n",
    "            predictions = torch.cat(predictions, 0)\n",
    "            ratings = torch.cat(ratings, 0)\n",
    "    print('RMSE:', rmse(predictions, ratings).item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Optional) Defining Minibatch Sampler for Implicit Feedback\n",
    "\n",
    "For implicit feedback datasets, instead of predicting the rating between a user and an item, we often need to predict whether a user would interact the item.  Since there is no rating for us to compute the loss, we instead rely on negative sampling for training.\n",
    "\n",
    "An example of training with negative sampling goes as follows.  For each user-item pair $(i,j)$ that appeared as an interaction (called *positive* example), we corrupt the item with a randomly-sampled item $j'$ (called *negative* example).  The model is then responsible for distinguishing whether the user-item pair is positive or negative.  The loss function for training implicit feedback can typically be any binary classification loss, for instance, a hinge loss:\n",
    "\n",
    "$$\n",
    "\\mathcal{L} = \\sum_{(i,j)\\in D, j'} \\max(\\hat{r}_{i,j'} - \\hat{r}_{i,j} + 1, 0)\n",
    "$$\n",
    "\n",
    "In graph learning we usually refer to this problem as *link prediction*, i.e. predicting whether an edge exists between two nodes.  In DGL it could be achieved with the following steps (the difference from rating prediction highlighted in bold):\n",
    "\n",
    "1. **Create two *pair graphs* that consists of all the users and items in the dataset, but only with the interactions in the training minibatch of positive examples and negative examples respectively.**\n",
    "2. *Compact* **both graphs** to remove all unnecessary nodes, keeping only the nodes with some edges going in or out.\n",
    "3. Construct blocks bottom-up for computing the multi-layer output of the nodes, as in the previous tutorial.\n",
    "4. Propagate messages top-down to obtain the output from the GNN.\n",
    "5. Copy the output to **both pair graphs** and compute the preference with `apply_edges()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinkPredictionMinibatchSampler(MinibatchSampler):\n",
    "    def __init__(self, graph, num_layers):\n",
    "        self.graph = graph\n",
    "        self.num_layers = num_layers\n",
    "        \n",
    "    def sample(self, batch):\n",
    "        # Convert the list of user-item-rating triplets into a pairs of users and items\n",
    "        users, items, _ = zip(*batch)\n",
    "        users = torch.stack(users)\n",
    "        items = torch.stack(items)\n",
    "        # Get corrupted items for negative examples\n",
    "        neg_items = torch.randint(0, self.graph.number_of_nodes('item'), (len(users),))\n",
    "        \n",
    "        # Create a pair graph for positive examples and negative examples (Step 1)\n",
    "        pos_pair_graph = dgl.heterograph(\n",
    "            {('user', 'watched', 'item'): (users, items)},\n",
    "            num_nodes_dict={'user': self.graph.number_of_nodes('user'), 'item': self.graph.number_of_nodes('item')})\n",
    "        neg_pair_graph = dgl.heterograph(\n",
    "            {('user', 'watched', 'item'): (users, neg_items)},\n",
    "            num_nodes_dict={'user': self.graph.number_of_nodes('user'), 'item': self.graph.number_of_nodes('item')})\n",
    "        \n",
    "        # Compact the graph (Step 2)\n",
    "        pos_pair_graph, neg_pair_graph = dgl.compact_graphs([pos_pair_graph, neg_pair_graph])\n",
    "        \n",
    "        # Construct blocks (Step 3)\n",
    "        # Note that pos_pair_graph and neg_pair_graph have the same set of users and items, so we only need\n",
    "        # to check one of them to get the seed nodes.\n",
    "        seeds = {'user': pos_pair_graph.nodes['user'].data[dgl.NID],\n",
    "                 'item': pos_pair_graph.nodes['item'].data[dgl.NID]}\n",
    "        # Note that here we would also remove edges connecting between users and both the corresponding positive\n",
    "        # and negative items appearing in the minibatch.\n",
    "        blocks = self.construct_blocks(seeds, (torch.cat([users, users]), torch.cat([items, neg_items])))\n",
    "        \n",
    "        # Copy node features from original graph to the sampled block.\n",
    "        # Note that for our model we only need to copy the features to the source side of the first block.\n",
    "        for feature_name in self.graph.nodes['user'].data.keys():\n",
    "            blocks[0].srcnodes['user'].data[feature_name] = \\\n",
    "                self.graph.nodes['user'].data[feature_name][blocks[0].srcnodes['user'].data[dgl.NID]]\n",
    "        for feature_name in self.graph.nodes['item'].data.keys():\n",
    "            blocks[0].srcnodes['item'].data[feature_name] = \\\n",
    "                self.graph.nodes['item'].data[feature_name][blocks[0].srcnodes['item'].data[dgl.NID]]\n",
    "            \n",
    "        return pos_pair_graph, neg_pair_graph, blocks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With minibatch sampler defined for link prediction, we only need to rewrite the training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
